from typing import Union
import uvicorn
from fastapi import FastAPI, Request, HTTPException
from app.api import *
from app.api import file
from fastapi.middleware.httpsredirect import HTTPSRedirectMiddleware
from fastapi.middleware.trustedhost import TrustedHostMiddleware
from fastapi.middleware.cors import CORSMiddleware
import asyncio
from contextlib import asynccontextmanager
from fastapi.staticfiles import StaticFiles
import os

@asynccontextmanager
async def lifespan(app: FastAPI):
    """处理应用程序的生命周期事件"""
    # 启动事件
    from app.api.ws import ws_handler
    from app.services.message import MessageService
    
    # 初始化WebSocket处理器
    await ws_handler.setup()
    print("WebSocket服务已初始化")
    
    # 初始化消息重试任务
    await MessageService.setup_retry_task()
    print("消息重试服务已初始化")
    
    yield  # 应用程序运行
    
    # 关闭事件
    for user_id in list(ws_handler.active_connections.keys()):
        try:
            await ws_handler.disconnect(user_id)
        except:
            pass
    print("WebSocket服务已关闭")
    
    # 取消消息重试任务
    if MessageService._retry_task:
        MessageService._retry_task.cancel()
        print("消息重试服务已关闭")

app = FastAPI(lifespan=lifespan)
# app.add_middleware(HTTPSRedirectMiddleware)  # 强制HTTPS
app.add_middleware(
    TrustedHostMiddleware,
    allowed_hosts=["127.0.0.1", "localhost"]  # 限制允许的host
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
app.include_router(user.router)
app.include_router(friend.router)
app.include_router(conversation.router)
app.include_router(message.router)
app.include_router(room.router)
app.include_router(other.router)
app.include_router(auth.router)
app.include_router(ws.router)
app.include_router(file.router)

# 注册静态资源目录
if not os.path.exists('uploads/images'):
    os.makedirs('uploads/images')
if not os.path.exists('uploads/files'):
    os.makedirs('uploads/files')
app.mount("/static/images", StaticFiles(directory="uploads/images"), name="images")
app.mount("/static/files", StaticFiles(directory="uploads/files"), name="files")

@app.get("/")
def read_root():
    return {"Hello": "World"}


@app.get("/items/{item_id}")
def read_item(item_id: int, q: Union[str, None] = None):
    return {"item_id": item_id, "q": q}

# @app.middleware("http")
# async def add_security_headers(request: Request, call_next):
#     response = await call_next(request)
#     # 添加安全相关的HTTP头
#     response.headers["X-Frame-Options"] = "DENY"
#     response.headers["X-Content-Type-Options"] = "nosniff"
#     response.headers["X-XSS-Protection"] = "1; mode=block"
#     response.headers["Content-Security-Policy"] = "default-src 'self'"
#     return response

if __name__ == "__main__":
    uvicorn.run(
        "app.main:app",
        host="0.0.0.0",
        port=8831,
        reload=True,  # 开发模式下启用热重载
        ws_ping_interval=30,  # WebSocket ping 间隔缩短到20秒
        ws_ping_timeout=30,   # WebSocket ping 超时也缩短到20秒
    )